import dns.nameserver
import dns.exception
import rtoml
import pandas as pd
import dns.message
import dns.rdataclass
import dns.rdatatype
import dns.query
import httpx


def str2server(dns_str: str):
    nameserver = dns_str
    if dns_str.startswith("https"):
        # nameserver = dns.nameserver.DoHNameserver(dns_str)
        query_func = dns.query.https
    elif dns_str.startswith("tls"):
        # nameserver = dns.nameserver.DoTNameserver(dns_str[6:])
        nameserver = dns_str[6:]
        query_func = dns.query.tls
    else:
        # nameserver = dns.nameserver.Do53Nameserver(dns_str)
        query_func = dns.query.udp
    return nameserver, query_func


def main():
    dns_list_file = "all_dns_list.toml"
    dst_path = "valid_dns_list.csv"
    record_max = 100
    timeout = 1
    with open(dns_list_file, 'r', encoding='utf-8') as fr:
        dns_list = rtoml.load(fr)

    valid_dns_pd = pd.DataFrame(columns=["DNS地址", "区域", "DNS描述", "解析结果", "响应时间"])
    area_list = list(dns_list.keys())
    area_list = ['国内', '港台', '海外'] # , '电信', '联通', '移动']

    qname = dns.name.from_text("baidu.com")
    q = dns.message.make_query(qname, dns.rdatatype.A)

    record_idx = 0
    for area in area_list:
        dns_dict = dns_list[area]
        for key in dns_dict:
            if record_idx > record_max:
                break
            dns_server, query_func = str2server(key)
            try:
                r = query_func(q, dns_server, timeout=timeout)
            except dns.exception.Timeout:
                print(key, "Timeout")
            except dns.exception.SyntaxError:
                print(key, "SyntaxError")
            except dns.exception.DNSException:
                print(key, "DNSException")
            except ValueError:
                print(key, "ValueError")
            except OSError:
                print(key, "OSError")
            except httpx.ConnectError:
                print(key, "ConnectError")
            except ConnectionResetError:
                print(key, "ConnectionResetError")
            else:
                record_idx += 1
                ns_rrset = r.get_rrset(r.answer, qname, dns.rdataclass.IN, dns.rdatatype.A)
                if ns_rrset is None:
                    print(key, "None")
                    continue
                rdata = list(ns_rrset.items.keys())[0]
                ans_ip = rdata.to_text()
                res_time = int(r.time * 1000)
                print(key, ans_ip, res_time)
                valid_dns_pd.loc[valid_dns_pd.shape[0]] = [key, area, dns_dict[key], ans_ip, res_time]

    valid_dns_pd.to_csv(dst_path, index=False, encoding='gbk')
    return 0


if __name__ == '__main__':
    main()
